# **************************************
# --*-- coding: utf-8 --*--
# @Author  : white
# @FileName: main.py
# @Time    : 2025-09-03
# **************************************
import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['axes.unicode_minus'] = False  # 解决坐标轴负号显示问题


# KD-tree节点类
class KDNode:
    """KD-tree的节点类"""

    def __init__(self, point=None, label=None, left=None, right=None, axis=None):
        self.point = point  # 数据点坐标
        self.label = label  # 数据点标签
        self.left = left  # 左子树
        self.right = right  # 右子树
        self.axis = axis  # 分割轴


# KD-tree类
class KDTree:
    """KD-tree数据结构实现"""

    def __init__(self):
        self.root = None

    def build_tree(self, points, labels, depth=0):
        """构建KD-tree
        Args:
            points: 数据点数组
            labels: 标签数组
            depth: 当前深度
        Returns:
            KDNode: 构建的树节点
        """
        if len(points) == 0:
            return None

        # 确定分割维度（轮流使用各个维度）
        k = points.shape[1]  # 数据维度
        axis = depth % k

        # 按照当前维度排序
        sorted_indices = np.argsort(points[:, axis])
        points = points[sorted_indices]
        labels = labels[sorted_indices]

        # 选择中位数作为分割点
        median = len(points) // 2

        # 创建节点
        node = KDNode(
            point=points[median],
            label=labels[median],
            axis=axis
        )

        # 递归构建左右子树
        node.left = self.build_tree(points[:median], labels[:median], depth + 1)
        node.right = self.build_tree(points[median + 1:], labels[median + 1:], depth + 1)

        return node

    def insert(self, points, labels):
        """插入数据点并构建树"""
        self.root = self.build_tree(points, labels)

    def distance(self, p1, p2):
        """计算欧几里得距离"""
        return np.sqrt(np.sum((p1 - p2) ** 2))

    def search_knn(self, target, k):
        """KNN搜索
        Args:
            target: 目标点
            k: 最近邻数量
        Returns:
            list: k个最近邻的标签
        """
        if self.root is None:
            return []

        # 存储最近邻结果的列表
        best_neighbors = []

        def dfs_search(node, target, depth=0):
            """深度优先搜索最近邻"""
            if node is None:
                return

            # 计算当前节点到目标点的距离
            dist = self.distance(node.point, target)

            # 更新最近邻列表
            if len(best_neighbors) < k:
                best_neighbors.append((dist, node.label, node.point))
                best_neighbors.sort(key=lambda x: x[0])  # 按距离排序
            elif dist < best_neighbors[-1][0]:
                best_neighbors[-1] = (dist, node.label, node.point)
                best_neighbors.sort(key=lambda x: x[0])

            # 确定搜索方向
            axis = node.axis
            if target[axis] < node.point[axis]:
                near_subtree = node.left
                far_subtree = node.right
            else:
                near_subtree = node.right
                far_subtree = node.left

            # 搜索近侧子树
            dfs_search(near_subtree, target, depth + 1)

            # 判断是否需要搜索远侧子树
            # 存储量未达到最大数量，目标的轴的距离与切割点的轴的距离小于
            if (len(best_neighbors) < k) or (abs(target[axis] - node.point[axis]) < best_neighbors[-1][0]):
                dfs_search(far_subtree, target, depth + 1)

        # 开始搜索
        dfs_search(self.root, target)

        # 返回k个最近邻的标签
        return [neighbor[1] for neighbor in best_neighbors]


# KNN分类器类
class KNNClassifier:
    """KNN分类器"""

    def __init__(self, k=3):
        self.k = k
        self.kdtree = KDTree()

    def fit(self, X, y):
        """训练模型（构建KD-tree）"""
        self.kdtree.insert(X, y)

    def predict(self, X):
        """预测"""
        predictions = []
        for point in X:
            # 找到k个最近邻
            neighbors = self.kdtree.search_knn(point, self.k)

            if len(neighbors) == 0:
                predictions.append(0)  # 默认类别
                continue

            # 投票决定类别
            unique_labels, counts = np.unique(neighbors, return_counts=True)
            predicted_label = unique_labels[np.argmax(counts)]
            predictions.append(predicted_label)

        return np.array(predictions)


# 数据生成函数
def generate_sample_data():
    """生成3个类别的样本数据，每类100条"""
    np.random.seed(42)  # 设置随机种子确保结果可重复

    # 类别0：围绕(2, 2)的正态分布
    class_0 = np.random.multivariate_normal([2, 2], [[1, 0.5], [0.5, 1]], 100)
    labels_0 = np.zeros(100)

    # 类别1：围绕(6, 6)的正态分布
    class_1 = np.random.multivariate_normal([6, 6], [[1, -0.3], [-0.3, 1]], 100)
    labels_1 = np.ones(100)

    # 类别2：围绕(2, 6)的正态分布
    class_2 = np.random.multivariate_normal([2, 6], [[1.2, 0.2], [0.2, 0.8]], 100)
    labels_2 = np.full(100, 2)

    # 合并所有数据
    X = np.vstack([class_0, class_1, class_2])
    y = np.hstack([labels_0, labels_1, labels_2])

    return X, y


# 计算准确率
def calculate_accuracy(y_true, y_pred):
    """计算预测准确率"""
    return np.mean(y_true == y_pred)


# 数据可视化
def visualize_data(X, y, title="数据分布"):
    """可视化数据分布"""
    plt.figure(figsize=(10, 8))
    colors = ['red', 'blue', 'green']
    labels = ['类别0', '类别1', '类别2']

    for i in range(3):
        mask = y == i
        plt.scatter(X[mask, 0], X[mask, 1], c=colors[i], label=labels[i], alpha=0.7)

    plt.xlabel('特征1')
    plt.ylabel('特征2')
    plt.title(title)
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()


# 主函数
def main():
    print("=" * 50)
    print("KD-tree与KNN算法演示")
    print("=" * 50)

    # 1. 生成数据
    print("1. 生成训练数据...")
    X_train, y_train = generate_sample_data()
    print(f"训练数据形状: {X_train.shape}")
    print(f"标签分布: {np.bincount(y_train.astype(int))}")

    # 2. 创建测试数据
    print("\n2. 生成测试数据...")
    np.random.seed(123)
    X_test, y_test = generate_sample_data()
    X_test = X_test[:30]  # 使用前30个点作为测试
    y_test = y_test[:30]

    # 3. 训练KNN分类器
    print("\n3. 构建KD-tree并训练KNN分类器...")
    knn = KNNClassifier(k=5)
    knn.fit(X_train, y_train)
    print("KD-tree构建完成!")

    # 4. 进行预测
    print("\n4. 进行预测...")
    y_pred = knn.predict(X_test)

    # 5. 计算准确率
    accuracy = calculate_accuracy(y_test, y_pred)
    print(f"\n5. 预测结果:")
    print(f"测试准确率: {accuracy:.3f}")

    # 6. 显示部分预测结果
    print(f"\n6. 前10个测试样本的预测结果:")
    print("真实标签:", y_test[:10].astype(int))
    print("预测标签:", y_pred[:10].astype(int))

    # 7. 可视化结果
    print("\n7. 数据可视化...")
    visualize_data(X_train, y_train, "训练数据分布")

    print("\n算法执行完成!")


if __name__ == "__main__":
    main()
